import torch
from torch.utils.data import DataLoader
from transformers import BertForTokenClassification
from data_processer import generate_datasets, create_mini_batch

model = BertForTokenClassification.from_pretrained("bert-base-uncased", num_labels=14)
model.train()
s = generate_datasets()
train_loader = DataLoader(s.train_dataset, batch_size=16, collate_fn=create_mini_batch, drop_last=False)
test_loader = DataLoader(s.test_dataset, batch_size=16, collate_fn=create_mini_batch, drop_last=False)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-5)
Epochs = 10
for epoch in range(Epochs):
    losses = 0.0
    for step, data in enumerate(train_loader):
        tokens_tensors, masks_tensors, label_tensors = [t for t in data]
        optimizer.zero_grad()
        outputs = model(input_ids=tokens_tensors, attention_mask=masks_tensors, labels=label_tensors)
        loss = outputs[0]
        loss.backward()
        optimizer.step()
        losses += loss.item()
        rate = (step + 1) / len(train_loader)
        a = "*" * int(rate * 50)
        b = "." * int((1 - rate) * 50)
        print("\rtrain loss: {:^3.0f}%[{}->{}]{:.3f}".format(int(rate * 100), a, b, loss), end="")
    print(losses)
